from sklearn.model_selection import GridSearchCV
from sklearn.neural_network import MLPClassifier
from sklearn.model_selection import train_test_split
from joblib import dump
import logging
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.metrics import precision_score, classification_report, recall_score, f1_score, confusion_matrix
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import numpy as np

class BaseModle:
    def __init__(self,x,y):
        self.x = x
        self.y = y
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(x, y, test_size=0.2, random_state=42)
        logging.basicConfig(filename='./classfication.log', level=logging.DEBUG)

    def mlp(self):
        ## 这个超级慢
        mlp = MLPClassifier(hidden_layer_sizes=(2048, 100), activation='relu',
                        validation_fraction=0.1, solver='adam',
                        alpha=0.0001, learning_rate_init=0.001)
        mlp.fit(self.X_train, self.y_train)
        self.y_pred = mlp.predict(self.X_test)
        logging.info('precision_score', precision_score(self.y_test, self.y_pred, average='weighted'))
        logging.info('recall_score', recall_score(self.y_test, self.y_pred, average='weighted'))
        logging.info('f1_score', f1_score(self.y_test, self.y_pred, average='weighted'))
        logging.info(classification_report(self.y_test, self.y_pred))
        # plt.plot(mlp.loss_curve_)
        dump(mlp, 'mlp2.joblib')


    def clf(self):
        clf = ExtraTreesClassifier()
        clf.fit(self.X_train,self.y_train)
        self.y_pred = clf.predict(self.X_test)
        logging.info('precision_score', precision_score(self.y_test, self.y_pred, average='weighted'))
        logging.info('recall_score', recall_score(self.y_test, self.y_pred, average='weighted'))
        logging.info('f1_score', f1_score(self.y_test, self.y_pred, average='weighted'))
        dump(clf, 'clf2.joblib')

    def torch(self):
        myNet = nn.Sequential(
                nn.Linear(2215, 100),
                nn.ReLU(),
                nn.Linear(100, 32),
                nn.Sigmoid()
            )
        optimzer = torch.optim.SGD(myNet.parameters(), lr=0.05)
        loss_func = nn.MSELoss()

        x = torch.tensor(self.x).float()
        y = torch.tensor(self.y).float()
        for epoch in range(5): ## 设置过大跑不动=-=
            out = myNet(x)
            loss = loss_func(out, y)  # 计算误差
            optimzer.zero_grad()  # 清除梯度
            loss.backward()
            optimzer.step()
            print(epoch)
        print(myNet(x).data) 
    
    def export_img(self):
        model_conf = confusion_matrix(np.argmax(self.y_test, axis=1), np.argmax(self.y_pred, axis=1))
        plt.imshow(model_conf, interpolation="nearest", cmap=plt.cm.Blues)
        plt.colorbar()
        tick_marks = np.arange(32)
        # plt.xticks(tick_marks, n_classes, rotation=45)
        # plt.yticks(tick_marks, n_classes)
        plt.xlabel("Actual Species")
        plt.ylabel("Predicted Species")
        plt.show()